import pandas as pd
import numpy as np
import sys
import warnings
warnings.filterwarnings('ignore')

train = pd.read_hdf('./input/train.h5')
test = pd.read_hdf('./input/test.h5')


if len(sys.argv) > 1:
    n = int(sys.argv[1])
else:
    n = 100
num_per = n

# train_1 = train.drop_duplicates(['ship'])[['ship', 'type']].reset_index().drop('index', axis=1)
# for q in np.linspace(0, 1, num_per):
#     def per(arr):
#         return np.quantile(arr, q)
#     # print(q, ':')
#     # 速度分级
#     for fea in ['v', 'x', 'y', 'd']:
#         t = train.groupby('ship')[f'{fea}'].agg({f'{fea}_{q}': per}).reset_index()
#         train_1 = pd.merge(train_1, t, on='ship', how='right')
#
# type_map = {'围网': 0, '拖网': 1, '刺网': 2}
# type_map_rev = {v: k for k, v in type_map.items()}
# train_1['type'] = train_1['type'].map(type_map)
# features = [x for x in train_1.columns if x not in ['ship', 'type']] + ['ship', 'type']
# train_1[features].to_csv(f'data/trn_per_{num_per}.csv', header=True, index=False)


train = test
if len(sys.argv) > 1:
    n = int(sys.argv[1])
else:
    n = 100
num_per = n
train_1 = train.drop_duplicates(['ship'])[['ship']].reset_index().drop('index', axis=1)
for q in np.linspace(0, 1, num_per):
    def per(arr):
        return np.quantile(arr, q)
    # print(q, ':')
    # 速度分级
    for fea in ['v', 'x', 'y', 'd']:
        t = train.groupby('ship')[f'{fea}'].agg({f'{fea}_{q}': per}).reset_index()
        train_1 = pd.merge(train_1, t, on='ship', how='right')

features = [x for x in train_1.columns if x not in ['ship']] + ['ship']
train_1[features].to_csv(f'data/tst_per_{num_per}.csv', header=True, index=False)